import PIL
import cv2
import numpy as np
from os import listdir
from os.path import join
from labelme.utils.draw import label_colormap


if __name__ == '__main__':
    root = 'orgin'
    types = ['train', 'val']
    idx = 0
    for typ in types:
        path = join(root, typ)
        img1_paths = listdir(join(path, 'image1'))
        img2_paths = listdir(join(path, 'image2'))
        lab_paths = listdir(join(path, 'label'))
        with open('VOCdevkit/VOC2007/ImageSets/Segmentation/{}.txt'.format(typ), 'w') as f:
            for img1, img2, lab in zip(img1_paths, img2_paths, lab_paths):
                img1 = join(path, 'image1', img1)
                img2 = join(path, 'image2', img2)
                lab = join(path, 'label', lab)
                img1 = cv2.imread(img1, 3)
                img2 = cv2.imread(img2, 3)
                lab = cv2.imread(lab, 0) / 255.

                cv2.imwrite('VOCdevkit/VOC2007/JPEGImages/{}.jpg'.format(idx), img1)
                cv2.imwrite('VOCdevkit/VOC2007/JPEGImages2/{}.jpg'.format(idx), img2)

                lbl_pil = PIL.Image.fromarray(lab.astype(np.uint8), mode='P')
                colormap = label_colormap(255)
                lbl_pil.putpalette((colormap * 255).astype(np.uint8).flatten())
                lbl_pil.save('VOCdevkit/VOC2007/SegmentationClass/{}.png'.format(idx))
                f.writelines('{}\n'.format(idx))
                print(idx)
                idx += 1


